package uratelimit

import (
	"golang.org/x/time/rate"
	"fmt"
)

// GetTokenBucket 获取一个bucket
func GetTokenBucket(key string) (*rate.Limiter, error) {
	bucket, ok := tokenBuckets[key]
	if ok != true {
		return nil, fmt.Errorf("no bucket: %s", key)
	}

	return bucket, nil
}

// RegisterTokenBucket 注册一个bucket
func RegisterTokenBucket(key string, r rate.Limit, b int) error {
	if _, exist := tokenBuckets[key]; exist == true {
		return nil
	}

	bucket := rate.NewLimiter(r, b)
	tokenBuckets[key] = bucket

	return nil
}

// ChangeTokenBucket 修改token bucket的配置参数
func ChangeTokenBucket(key string, r rate.Limit, b int) error {
	bucket, ok := tokenBuckets[key]
	if ok != true {
		return fmt.Errorf("no bucket: %s", key)
	}

	// 重新申请一个
	bucket = rate.NewLimiter(r, b)
	tokenBuckets[key] = bucket
	return nil
}

var tokenBuckets map[string]*rate.Limiter

func init()  {
	tokenBuckets = make(map[string]*rate.Limiter)
}